package com.czkevin.biliwsclient;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import lombok.Getter;
import org.brotli.dec.BrotliInputStream;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.zip.InflaterInputStream;

public class BiliPacket {

    public final long time = System.currentTimeMillis();

    public int action;
    public int param;

    @Getter
    public byte[] data;

    public JSONObject parsedJSON = null;

    public BiliPacket() {
    }

    public BiliPacket(int action, int param) {
        this.action = action;
        this.param = param;
    }

    public BiliPacket(int action, int param, byte[] data) {
        this.action = action;
        this.param = param;
        this.data = data;
    }

    public BiliPacket(int action, int param, String data) {
        this.action = action;
        this.param = param;
        this.data = data.getBytes(StandardCharsets.UTF_8);
    }

    public byte[] serialize() {
        try {
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            DataOutputStream dos = new DataOutputStream(bos);

            int data_size = data != null ? data.length : 0;
            int size = 16 + data_size;

            dos.writeInt(size);
            dos.writeShort(16); // magic?
            dos.writeShort(1); // version
            dos.writeInt(action);
            dos.writeInt(param);

            if(data != null) dos.write(data);

            dos.close();
            bos.close();
            return bos.toByteArray();
        } catch (Exception e){
            e.printStackTrace();
            return null;
        }
    }

    public BiliPacket setAction(int action) {
        this.action = action;
        return this;
    }

    public BiliPacket setParam(int param) {
        this.param = param;
        return this;
    }

    public BiliPacket setData(byte[] data) {
        this.data = data;
        return this;
    }

    public BiliPacket setPacket(String data) {
        this.data = data.getBytes(StandardCharsets.UTF_8);
        return this;
    }

    public JSONObject getDataAsJSON() {
        if(parsedJSON != null) return parsedJSON;
        if(data == null) return null;
        String str = new String(data, StandardCharsets.UTF_8);
        parsedJSON = JSON.parseObject(str);
        return parsedJSON;
    }

    public static BiliPacket[] deserialize(byte[] data) {
        return deserialize(data, false);
    }
    public static BiliPacket[] deserialize(byte[] data, boolean inside) {
        try {
            ByteArrayInputStream bis = new ByteArrayInputStream(data);
            DataInputStream dis = new DataInputStream(bis);

            List<BiliPacket> packets = new ArrayList<>();
            try {
                while (dis.available() > 0) {
                    BiliPacket packet = new BiliPacket();

                    int total_len = dis.readInt();
                    // int magic = dis.readUnsignedShort();
                    dis.skipBytes(2); // magic
                    int ver = dis.readUnsignedShort();
                    int action = dis.readInt();
                    int param = dis.readInt();
                    byte[] d = null;

                    if (total_len - 16 > 0) {
                        d = new byte[total_len - 16];
                        dis.readFully(d);
                    }

                    /*System.out.println(String.format("len=%d, magic=%d, ver=%d, action=%d, param=%d, data=%s",
                            total_len, magic, ver, action, param, d == null ? "NULL" : d.length + "b"));*/

                    packet.action = action;
                    packet.param = param;
                    packet.data = d;

                    if(inside || ver == 1 || ver == 0) {
                        // Version 1: Straight ahead, no shitty stuff.
                        packets.add(packet);
                    } else if(ver == 2 || ver == 3) {
                        // Version 2: Deflated shit, now we do some more shit! >:(
                        // Version 3: Brotli shit, added in late 2022.
                        InputStream iis;
                        if (ver == 2) {
                            iis = new InflaterInputStream(new ByteArrayInputStream(d));
                        } else if (ver == 3) {
                            iis = new BrotliInputStream(new ByteArrayInputStream(d));
                        } else {
                            throw new IllegalStateException(String.format("unknown packet version %d", ver));
                        }
                        byte[] decompressed = readFully(iis);
                        iis.close();
                        BiliPacket[] hiddenPackets = deserialize(decompressed, true);
                        if(hiddenPackets != null) {
                            Collections.addAll(packets, hiddenPackets);
                        }
                    }
                }
            } finally {
                bis.close();
                dis.close();

                return packets.toArray(new BiliPacket[0]);
            }
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }

    private static byte[] readFully(InputStream is) throws IOException {
        ByteArrayOutputStream bos = new ByteArrayOutputStream();
        byte[] buff = new byte[2048];
        int read;
        while((read = is.read(buff)) != -1) {
            bos.write(buff, 0, read);
        }
        return bos.toByteArray();
    }
}
